package org.me.dubbodemo.appbase.Util;

import org.apache.dubbo.common.utils.StringUtils;
import org.apache.dubbo.rpc.RpcContext;
import org.me.dubbodemo.appbase.constants.TraceConst;
import org.slf4j.MDC;

import java.util.UUID;

public class TraceUtil {
    public static String getOrInitTraceFromMDC() {
        String trace = MDC.get(TraceConst.TRACE_KEY);
        if (StringUtils.isBlank(trace)) {
            trace = initTrace();
            MDC.put(TraceConst.TRACE_KEY, trace);
        }
        return trace;
    }

    public static String initTrace() {
        return UUID.randomUUID().toString();
    }

    public static String getOrInitTraceFromRpcContext() {
        String trace = RpcContext.getServerAttachment().getAttachment(TraceConst.TRACE_KEY);
        if (StringUtils.isBlank(trace)) {
            trace = initTrace();
        }
        return trace;
    }

    public static void populateTraceToMDC(final String trace) {
        MDC.put(TraceConst.TRACE_KEY, trace);
    }
}
